package com.coocaa.typhoon.api.third.volcengine.provider;

import cn.idev.excel.util.StringUtils;
import com.alibaba.fastjson.JSON;
import com.coocaa.cirrus.common.base.exception.BusinessException;
import com.coocaa.cirrus.common.third.filestore.config.TosProperties;
import com.coocaa.typhoon.api.third.volcengine.constant.EmbeddingTypeConstant;
import com.coocaa.typhoon.api.third.volcengine.model.EmbeddingQueryDTO;
import com.coocaa.typhoon.api.third.volcengine.model.VikingDbEmbeddingResponse;
import com.coocaa.typhoon.api.third.volcengine.util.OkhttpSign;
import lombok.extern.slf4j.Slf4j;
import okhttp3.ConnectionPool;
import okhttp3.OkHttpClient;
import org.apache.commons.collections4.CollectionUtils;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * @author bijiahao
 * @date : 2025/9/25 16:30
 * @description
 */
@Slf4j
@Service
public class VolcEmbeddingProvider implements InitializingBean {
    private OkHttpClient client;
    @Autowired
    private TosProperties tosProperties;

    @Override
    public void afterPropertiesSet() throws Exception {
        OkHttpClient.Builder clientBuilder = new OkHttpClient.Builder();
        //读取超时
        clientBuilder.readTimeout(3, TimeUnit.SECONDS);
        //连接超时
        clientBuilder.connectTimeout(5, TimeUnit.SECONDS);
        //写入超时
        clientBuilder.writeTimeout(3, TimeUnit.SECONDS);
        //自定义连接池最大空闲连接数和等待时间大小，否则默认最大5个空闲连接
        clientBuilder.connectionPool(new ConnectionPool(32, 5, TimeUnit.MINUTES));
        client = clientBuilder.retryOnConnectionFailure(true).build();
    }

    public List<Float> questionVector(EmbeddingQueryDTO queryDTO) {
        List<EmbeddingQueryDTO> queryDTOList = new ArrayList<>();
        queryDTOList.add(queryDTO);
        Map<String, List<Float>> questionVectorMap = embedding(queryDTOList);
        List<Float> floats = questionVectorMap.get(queryDTO.getContent());
        return floats;
    }


    public Map<String, List<Float>> embedding(List<EmbeddingQueryDTO> queryDTOList) {
        String path = "/api/vikingdb/embedding";
        OkhttpSign okhttpSign = new OkhttpSign(tosProperties.getTosRegionName(), path, tosProperties.getTosSecretId(), tosProperties.getTosSecretKey(), this.client);
        HashMap<String, String> queryMap = new HashMap<>();
        HashMap<String, Object> bodyMap = new HashMap<>();
        bodyMap.put("dense_model", buildDenseModel());

        bodyMap.put("data", buildData(queryDTOList));
        String json = JSON.toJSONString(bodyMap);
        String resStr = okhttpSign.doRequest("POST", queryMap, json);
        if (StringUtils.isBlank(resStr)) {
            return null;
        }
        VikingDbEmbeddingResponse response = JSON.parseObject(resStr, VikingDbEmbeddingResponse.class);
        if (response == null) {
            return null;
        }
        VikingDbEmbeddingResponse.EmbeddingResult result = response.getResult();
        if (result == null) {
            return null;
        }
        List<VikingDbEmbeddingResponse.DataItem> data = result.getData();
        if (CollectionUtils.isEmpty(data)) {
            return null;
        }
        int size = queryDTOList.size();
        Map<String, List<Float>> finalRes = new HashMap<>();
        for (int i = 0; i < size; i++) {
            EmbeddingQueryDTO queryDTO = queryDTOList.get(i);
            VikingDbEmbeddingResponse.DataItem dataItem = data.get(i);
            finalRes.put(queryDTO.getContent(), dataItem.getDense());
        }
        return finalRes;
    }


    private static List<HashMap<String, Object>> buildData(List<EmbeddingQueryDTO> queryDTOList) {
        List<HashMap<String, Object>> dataList = new ArrayList<>();
        for (EmbeddingQueryDTO embeddingQueryDTO : queryDTOList) {
            HashMap<String, Object> param = buildData(embeddingQueryDTO);
            dataList.add(param);
        }
        return dataList;
    }

    private static HashMap<String, Object> buildData(EmbeddingQueryDTO queryDTO) {
        //https://www.volcengine.com/docs/84313/1791161
        HashMap<String, Object> param = new HashMap<>();
        if (EmbeddingTypeConstant.TEXT.equals(queryDTO.getType())) {
            param.put(EmbeddingTypeConstant.TEXT, queryDTO.getContent());
        } else if (EmbeddingTypeConstant.IMAGE.equals(queryDTO.getType())) {
            String tosLink = queryDTO.getContent();
            param.put(EmbeddingTypeConstant.IMAGE, tosLink);
        } else if (EmbeddingTypeConstant.VIDEO.equals(queryDTO.getType())) {
            String tosLink = queryDTO.getContent();
            HashMap<String, Object> videoMap = new HashMap<>();
            videoMap.put("value", tosLink);
            videoMap.put("fps", 0.4f);
            param.put(EmbeddingTypeConstant.VIDEO, videoMap);
        } else {
            throw new BusinessException("未知类型:" + queryDTO.getType());
        }
        return param;
    }

    private static HashMap<String, Object> buildDenseModel() {
        HashMap<String, Object> paramMap = new HashMap<>();
        paramMap.put("name", "doubao-embedding-vision");
        paramMap.put("version", "250615");
        paramMap.put("dim", 2048);
        return paramMap;
    }
}
